using namespace std;

#include <iostream>
#include <fstream>
#include <sys/stat.h>
#include <time.h>
#include "SNP.h"
#include "utils.h"
#include "Sample.h"
#include <algorithm>
#include <vector>
#include <iomanip>
#include <string.h>
// threshold for vote
#define VOTE_THRESHOLD 0.2

typedef struct stat Stat;


static const string prefixFolder = "snpData/"; // object folder
static string inputFile = "input.txt"; // input file
static string outputFile = "output.txt"; // output file
static string snpListFile = "snpList.txt"; // stores list of snps
static float threshold = 0.1; // threshold for votes
static int numberOfSNPs = 0; // total number of SNPs
static vector<string> sampleNames; // list of samples names
static double currentLoglikelihood = 0; // the log likelihood that we want to maximize
static double currentLoglikelihood1 = 0;
static vector<string> removedList; // list of bad samples which have been removed
static string removedSample = "";
// this struct is used to store the votes of sample
static vector<string> snpNames; // list of snp names

struct BadSample {
	string name;
	int numOfVotes;
};
static vector<BadSample> badSamples; // list of potential bad samples


/**
 * add a candidate to list of bad samples
 *
 * find a sample in list which have the same name as given sample
 * if this sample exists, increase its votes by 1
 * else add a new sample with 1 vote
 */
 
void putToList(Sample s) {
	int index = -1;
	for(unsigned int i=0; i<badSamples.size(); i++) {
		if(badSamples[i].name == s.getName()) {
			index = i;
			break;
		}
	}
	if(index > -1) {
		badSamples[index].numOfVotes += 1;
	}
	else {
		BadSample tmp;
		tmp.name = s.getName();
		tmp.numOfVotes = 1;
		badSamples.push_back(tmp);
	}
}

/******************************************************************************************
 *	Function for test
 ******************************************************************************************/

static int do_mkdir(const char *path, mode_t mode)
{
    Stat st;
    int  status = 0;

    if (stat(path, &st) != 0)
    {
        /* Directory does not exist */
        if (mkdir(path, mode) != 0)
            status = -1;
    }
    else if (!S_ISDIR(st.st_mode))
    {
        status = -1;
    }

    return(status);
}

/**
 * given the input file, create objects and store them in prefix folder
 * get list of samples, list of candidate, list of snp, ...
 */
void processFromFile(string fileName) {
	ifstream myFile;
	ofstream snpSaveToFile;
	snpSaveToFile.open(snpListFile.c_str(), ios::out);
	myFile.open(fileName.c_str(), ios::in);
	if (!myFile) {
		cout << "could not open file " << fileName << endl;
		return;
	}
	if (!snpSaveToFile) {
		cout << "could not create file " << snpListFile << endl;
		return;
	}
	string line = "";

	// read each line of input file
	// lines start with '##' are comment lines
	// line starts with '#CHROM' stores name list of samples
	// other line, add to a SNP to get data
	// write the list of SNPs to file
	while(!myFile.eof()) {
		getline(myFile, line);
		if (line == "") continue;
		else if (line.compare(0, 2, "##") == 0) {
			continue;
		}
		else if (line.compare(0, 6, "#CHROM") == 0) {
			sampleNames = getHeader(line);
		}
		else {
			numberOfSNPs ++;
			SNP tmpSNP;
			tmpSNP.assignData(line, sampleNames);
			tmpSNP._saveToFile(prefixFolder + tmpSNP.getName());
			vector<Sample> badCandidates = tmpSNP.getBadSamples();
			for(unsigned int i=0; i<badCandidates.size(); i++) {
				putToList(badCandidates[i]);
			}
			snpNames.push_back(tmpSNP.getName());
			snpSaveToFile << tmpSNP.getName() << endl;
		}
	}
	myFile.close();
	snpSaveToFile.close();
}

/**
 * load list of SNPs name from file
 */
void loadObjectList() {
	ifstream myFile;
	myFile.open(snpListFile.c_str(), ios::in);
	if (!myFile) {
		cout << "could not open file " << myFile << endl;
		return;
	}

	string line = "";
	while(!myFile.eof()) {
		getline(myFile, line);
		if(line == "") continue;
		else {
			numberOfSNPs ++;
			snpNames.push_back(line);
			SNP tmpSNP;
			tmpSNP._loadFromFile(prefixFolder + line);
			vector<Sample> badCandidates = tmpSNP.getBadSamples();
			for(unsigned int i=0; i<badCandidates.size(); i++) {
				putToList(badCandidates[i]);
			}
		}
	}
	myFile.close();
}

/******************************************************************************************
 *	
 ******************************************************************************************/




/**
 *	thantvk52
 *	Read file to get number of SNPs
 */

int getNumberOfSNPs(string fileName) {
	ifstream myFile;
	myFile.open(fileName.c_str(), ios::in);
	if (!myFile) {
		cout << "could not open file " << fileName << endl;
		return 0;
	}
	string line = "";

	// read each line of input file
	// lines start with '##' are comment lines
	// line starts with '#CHROM' stores name list of samples
	// other line, add to a SNP to get data
	int i = 0;
	while(!myFile.eof()) 
	{
		getline(myFile, line);
		if (line == "" || line.compare(0, 2, "##") == 0 || line.compare(0, 6, "#CHROM") == 0) 
			continue;
		else
			i++;
	}
	return i;
}

/**
 *	thantvk52
 *	Read file to get information of SNPs, save to an array
 */

void processFromFile1(string fileName, SNP snps[]) 
{
	ifstream myFile;
	myFile.open(fileName.c_str(), ios::in);
	if (!myFile) {
		cout << "could not open file " << fileName << endl;
		return;
	}
	string line = "";

	// read each line of input file
	// lines start with '##' are comment lines
	// line starts with '#CHROM' stores name list of samples
	// other line, add to a SNP to get data

	int i = 0;	
	while(!myFile.eof()) 
	{
		getline(myFile, line);
		if (line == "") continue;
		else if (line.compare(0, 2, "##") == 0) {
			continue;
		}
		else if (line.compare(0, 6, "#CHROM") == 0) {
			sampleNames = getHeader(line);
		}
		else {
			snps[i].assignData(line, sampleNames);
			vector<Sample> badCandidates = snps[i].getBadSamples();
			for(unsigned int j=0; j<badCandidates.size(); j++) {
				putToList(badCandidates[j]);
			}
			i++;
		}
	}	
	myFile.close();
}

/**
 * this function is used to compare two candidate by votes
 */
bool myCompareCandidate(BadSample a, BadSample b) {
	return a.numOfVotes > b.numOfVotes;
}

/**
 * sort the list of candidates by number of votes and only output samples having more than threshold votes
 */
vector<BadSample> getBadSamplesWithThreshold(float threshold = VOTE_THRESHOLD) {
	float thresholdSam = threshold * numberOfSNPs;
	sort(badSamples.begin(), badSamples.end(), myCompareCandidate);
	vector<BadSample> returnVec;
	for(unsigned int i=0; i<badSamples.size(); i++) {
		if(badSamples[i].numOfVotes < thresholdSam) {
			break;
		}
		returnVec.push_back(badSamples[i]);
	}
	return returnVec;
}

/**
 *	get SNP information in each line of input file
 */
/*
void getSNPInfo(SNP snps[]) {

	if(numberOfSNPs != 0) {
		
		for(int i = 0; i < numberOfSNPs; i++) {
			snps[i].assignData(dataInLine.at(i), sampleNames);
			vector<Sample> badCandidates = snps[i].getBadSamples();
			//cout << "bad sample of SNP " << i << ": " << badCandidates.size() << endl;
			for(unsigned int i=0; i<badCandidates.size(); i++) {
				putToList(badCandidates[i]);
			}
			//cout << "Number bad sample: " << badSamples.size() << endl;
		}
	}	
}*/

void compareSNP(SNP snp1, SNP snp2) {
	cout << "Stay here" << endl;
	if(snp1.getName() != snp2.getName()) {
		cout << "Different name" << endl;
		return;
	}
	cout << "Stay here" << endl;
	if(snp1.getNumberOfSamples() != snp2.getNumberOfSamples())
	{	
		cout << "Different number of sample" << endl;
	}
	cout << "Stay here" << endl;
}


/**
 * calculate the log likelihood of given SNPs
 * 	for each SNP, load its data from file (each SNP is stored in one file)
 * 	remove all samples in the list
 * 	calculate the log likelihood from each SNP
 *
 * return the sum
 */
double calculateTheLogLikelihood(vector<string> reSamples, string reSample, SNP snps[]) {
	double thisLogLikelihood = 0;
	double thisLogLikelihood1 = 0;
	for(int i=0; i<numberOfSNPs; i++) {
		SNP tmp;
		tmp._loadFromFile(prefixFolder + snpNames[i]);
		tmp.removeSamples(reSamples);
		double t1 = tmp.calculateLogLikelihood();
		//cout << snpNames[i] << "\t" << t << endl;
		thisLogLikelihood += t1;
		
		if(reSample != "")
			snps[i].removeSample(reSample);
		//cout << "remove finish" << endl;
		double t2 = snps[i].calculateLogLikelihood();
		thisLogLikelihood1 += t2;
		//if(t1 != t2) cout << setprecision(30)<< t1  << " " << t2 << endl;
		
		if(t1 != t2) {
			//cout << "Error loglikelihood" << endl;
			//compareSNP(snps[i], tmp);
			//double a[tmp.getNumberOfSamples()][3];
			//double b[snps[i].getNumberOfSamples()][3];
			cout << tmp.getNumberOfSamples() << " " << snps[i].getNumberOfSamples() << endl;
			//tmp.compare(snps[i]);			
			//snps[i].display(b);
			//tmp.display(a);
			/*
			for(int j = 0; i < tmp.getNumberOfSamples(); j++) {
				for(int k = 0; k < 3; k++) {
					if(a[j][k] != b[j][k]) {
						cout << setprecision(10) << "not equal: " << a[j][k] << " " << b[j][k] << endl;
					}
					//else
						//	cout << "not equal: " << a[j][k] << " " << b[j][k] << endl;
				}
			}*/
		}/*
		else {
			double a[tmp.getNumberOfSamples()][3];
			double b[snps[i].getNumberOfSamples()][3];
			snps[i].display(b);
			tmp.display(a);
			for(int j = 0; i < tmp.getNumberOfSamples(); j++) {
				for(int k = 0; k < 3; k++) {
					if(a[j][k] != b[j][k]) {
						cout << "not equal: " << a[j][k] << " " << b[j][k] << endl;
					}
					//else
						//	cout << "not equal: " << a[j][k] << " " << b[j][k] << endl;
				}
			}
		}*/
	}
	return thisLogLikelihood1;
}


/**
 *	author: thanhtv
 *	return loglikelihood
 */

double RemoveAndCalculateTheLogLikelihood(string reSample, SNP snps[]) {
	double thisLogLikelihood = 0;
	for(int i=0; i<numberOfSNPs; i++) {
		if(reSample != "")
			snps[i].removeSample(reSample);
		//cout << "remove finish" << endl;
		thisLogLikelihood += snps[i].calculateLogLikelihood();		
	}
	return thisLogLikelihood;
}

/**
 *	author: thanhtv
 */

void undo(SNP snps[]) {
	for(int i = 0; i < numberOfSNPs; i++) {
		snps[i].reAddSample();
	}
}
/*
void display(SNP snps[]) {
	for(int i = 0; i < numberOfSNPs; i++) {
		cout << "Data of snp number " << i << endl;
		snps[i].display();
	}
}


void display() {
	for(int i=0; i<numberOfSNPs; i++) {
		SNP tmp;
		tmp._loadFromFile(prefixFolder + snpNames[i]);
		tmp.display();
	}
}

*/
/**
 * the process which checks and removes bad samples
 *
 * 	calculate the log likelihood when the removedList is NULL, call it the last log likelihood
 *
 * 	for each candidate, create a new vector from removedList
 * 	add the candidate to the vector
 * 	calculate the new log likelihood
 *
 * 	if the new log likelihood is greater than the last log likelihood
 * 		update the last log likelihood
 * 		add the candidate to the removedList
 * 	else
 * 		check another candidates
 *
 * dump removedList to file
 */
void removeBadSamples(SNP snps[]) {
	// calculate the current log likelihood
	/**
	 * test ...
	 */
	//for (int i=0; i<230; i++) {
	//	removedList.push_back(badSamples[i].name);
	//}
	time_t start = clock();
	currentLoglikelihood = calculateTheLogLikelihood(removedList, "", snps);
	time_t time = clock() - start;
	cout << "shortest time to calculate the fitness: " << (double)time / ((double)CLOCKS_PER_SEC) << "s" << endl << endl << endl;
	//return;
	// format print
	cout << left << setw(40) << "Sample names" << "  " << setw(10) << "Votes" << "  " << setw(20)
			<< "Fitness before" << "  " << setw(20) << "Fitness after" << "  " << setw(7)
			<< "Remove?" << endl;
	cout << setfill('-') << setw(105) << "" << setfill(' ') << endl;

	for(unsigned int i=0; i<badSamples.size(); i++) {
		//vector<string> currentRemovedList(removedList);
		//currentRemovedList.push_back(badSamples[i].name);
		removedList.push_back(badSamples[i].name);
		double thisLoglikelihood = calculateTheLogLikelihood(removedList, badSamples[i].name, snps);
		if(thisLoglikelihood < currentLoglikelihood) {
			removedList.pop_back();
			undo(snps);
			cout << left << setw(40) << badSamples[i].name << "  " << setw(10) << badSamples[i].numOfVotes
					<< "  " << setw(20) << currentLoglikelihood << "  " << setw(20) << thisLoglikelihood
					<< "  " << setw(7) << "" << endl;
		}
		else {
			
			//removedList.push_back(badSamples[i].name);
			cout << left << setw(40) << badSamples[i].name << "  " << setw(10) << badSamples[i].numOfVotes
					<< "  " << setw(20) << currentLoglikelihood << "  " << setw(20) << thisLoglikelihood
					<< "  " << setw(7) << "X" << endl;
			currentLoglikelihood = thisLoglikelihood;
		}
	}
	cout << setfill('-') << setw(105) << "" << endl;
	cout << endl << "There are " << removedList.size() << " samples will be removed" << endl;

	ofstream myFile;
	myFile.open(outputFile.c_str(), ios::out);
	if (!myFile) {
		cout << "could not create file " << outputFile << endl;
		for(unsigned int i=0; i<removedList.size(); i++) {
			cout << removedList[i] << endl;
		}
		return;
	}
	for (unsigned int i=0; i<removedList.size(); i++) {
		myFile << removedList[i] << endl;
	}
	myFile.close();
}
 
void removeBadSamples1(SNP snps[]) {

	time_t start = clock();
	currentLoglikelihood1 = RemoveAndCalculateTheLogLikelihood("", snps);
	//cout << "Current loglikelihood: " << currentLoglikelihood << endl;
	time_t time = clock() - start;
	cout << "shortest time to calculate the fitness: " << (double)time / ((double)CLOCKS_PER_SEC) << "s" << endl << endl << endl;
	//return;
	// format print
	cout << left << setw(40) << "Sample names" << "  " << setw(10) << "Votes" << "  " << setw(20)
			<< "Fitness before" << "  " << setw(20) << "Fitness after" << "  " << setw(7)
			<< "Remove?" << endl;
	cout << setfill('-') << setw(105) << "" << setfill(' ') << endl;

	for(unsigned int i=0; i<badSamples.size(); i++) {
		removedList.push_back(badSamples[i].name);
		double thisLoglikelihood = RemoveAndCalculateTheLogLikelihood(badSamples[i].name, snps);
		if(thisLoglikelihood < currentLoglikelihood1) {
			cout << "Before undo" << endl;			
			//display(snps);			
			undo(snps);
			cout << "After undo" << endl;
			//display(snps);
			removedList.pop_back();
			cout << left << setw(40) << badSamples[i].name << "  " << setw(10) << badSamples[i].numOfVotes
					<< "  " << setw(20) << currentLoglikelihood1 << "  " << setw(20) << thisLoglikelihood
					<< "  " << setw(7) << "" << endl;
		}
		else {
			//removedList.push_back(badSamples[i].name);
			cout << left << setw(40) << badSamples[i].name << "  " << setw(10) << badSamples[i].numOfVotes
					<< "  " << setw(20) << currentLoglikelihood1 << "  " << setw(20) << thisLoglikelihood
					<< "  " << setw(7) << "X" << endl;
			currentLoglikelihood1 = thisLoglikelihood;
		}
	}
	cout << setfill('-') << setw(105) << "" << endl;
	cout << endl << "There are " << removedList.size() << " samples will be removed" << endl;

	ofstream myFile;
	myFile.open(outputFile.c_str(), ios::out);
	if (!myFile) {
		cout << "could not create file " << outputFile << endl;
		for(unsigned int i=0; i<removedList.size(); i++) {
			cout << removedList[i] << endl;
		}
		return;
	}
	for (unsigned int i=0; i<removedList.size(); i++) {
		myFile << removedList[i] << endl;
	}
	myFile.close();
}

int main(int argc, char* argv[])
{	bool hasList = false;
	if (argc > 1) {
		try {
			if (argc % 2 == 0) throw "invalid parameters";
			for(int i=1; i<argc-1; i++) {
				if(strcmp(argv[i], "-in")==0) {
					inputFile = argv[i+1];
					i+=1;
				}
				else if(strcmp(argv[i], "-out") == 0) {
					outputFile = argv[i+1];
					i+=1;
				}
				else if(strcmp(argv[i], "-th") == 0) {
					threshold = atof(argv[i+1]);
					i+=1;
				}
				else if(strcmp(argv[i], "-obj") == 0) {
					snpListFile = argv[i+1];
					hasList = true;
					i+=1;
				}
			}
		} catch(...) {
			cout << "invalid parameters" << endl;
			cout << "try: " << argv[0] << " " << "[-in inputFile] [-out outputFile] [-th threshold] [-obj snpListFile]" << endl;
			return 0;
		}
	}
	
	cout << "input file:\t" << inputFile << endl;
	cout << "output file:\t" << outputFile << endl;
	cout << "threshold:\t" << threshold << endl;
	cout << "object list:\t" << snpListFile << endl;
	cout << "================================================" << endl;
	clock_t start, time;
	start = clock();

	
	
	//vector<BadSample> badSamples1, badSamples2;
	
	
	numberOfSNPs = getNumberOfSNPs(inputFile);
	SNP snps[numberOfSNPs];	
	processFromFile1(inputFile, snps);
	badSamples.clear();
	numberOfSNPs = 0;
	do_mkdir(prefixFolder.c_str(), 777); // create the object folder if not exists

	start = clock();
	if (!hasList) {
		processFromFile(inputFile);
	}
	else {
		loadObjectList();
	}
	time = clock() - start;

	cout << "total time to create objects from vcf file: " << (double)time / ((double)CLOCKS_PER_SEC) << "s" << endl;
	cout << "number of SNPs: " << numberOfSNPs << endl;

	start = clock();
	badSamples = getBadSamplesWithThreshold(threshold);
	cout << "number of candidates: " << badSamples.size() << endl;
	removeBadSamples(snps);
	time = clock() - start;

	cout << endl << endl << endl << "total time to remove bad samples: " << (double)time / ((double)CLOCKS_PER_SEC) << "s" << endl;
	cout << "================================================" << endl;
	cout << "removed samples are written into " << outputFile << endl;
	cout << "snps' names are written into " << snpListFile << endl;

	return 0;
}
